import torch.nn as nn
import torch
from torchvision import models
from .config import *

# 构建ResNet模型
class ResNetSpeaker(nn.Module):
    def __init__(self, model_type, num_classes=idAmount):
        super(ResNetSpeaker, self).__init__()
        # 加载预训练的ResNet模型
        if model_type == 'resnet18':
            self.resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        elif model_type == 'resnet34':
            self.resnet = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)
        elif model_type == 'resnet50':
            self.resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        elif model_type == 'resnet101':
            self.resnet = models.resnet101(weights=models.ResNet101_Weights.DEFAULT)
        else:
            raise KeyError('模型未定义')

        # 替换第一层卷积，以接受1个通道的输入
        self.resnet.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)

        # 调整avgpool的kernel_size
        self.resnet.avgpool = nn.AdaptiveAvgPool2d((1, 1))

        # 替换全连接层以匹配目标类别的数量
        self.resnet.fc = nn.Linear(self.resnet.fc.in_features, num_classes)

    def forward(self, x):
        x = self.resnet(x)
        return x
    

def conv3x3(in_planes, out_planes, stride=1, dilation=1):
    return nn.Conv2d(in_planes, out_planes, 3, stride=stride, padding=dilation, bias=False, dilation=dilation)

def conv1x1(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_channels, out_channels, stride=1, downsample=None, dilation=1, norm_layer=None):
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self.conv1 = conv3x3(in_channels, out_channels, stride, dilation)
        self.bn1 = norm_layer(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels, 1, dilation)
        self.bn2 = norm_layer(out_channels)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out
    

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_channels, out_channels, stride=1, downsample=None, dilation=1, norm_layer=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d

        self.conv1 = conv1x1(in_channels, out_channels)
        self.bn1 = norm_layer(out_channels)
        
        self.conv2 = conv3x3(out_channels, out_channels, stride, dilation)
        self.bn2 = norm_layer(out_channels)
        
        self.conv3 = conv1x1(out_channels, out_channels * self.expansion)
        self.bn3 = norm_layer(out_channels * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNet34Custom(nn.Module):
    def __init__(self, num_classes=idAmount):
        super(ResNet34Custom, self).__init__()
        self.inplanes = 32
        norm_layer = nn.BatchNorm2d

        # 为较小的灰度图设置合适的输入层
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = norm_layer(32)
        self.relu = nn.ReLU(inplace=True)

        # 设置模型各层
        self.layer1 = self._make_layer(BasicBlock, 32, 3, stride=1, norm_layer=norm_layer)
        self.layer2 = self._make_layer(BasicBlock, 64, 4, stride=2, norm_layer=norm_layer)
        self.layer3 = self._make_layer(BasicBlock, 128, 6, stride=2, norm_layer=norm_layer)
        self.layer4 = self._make_layer(BasicBlock, 256, 3, stride=2, norm_layer=norm_layer)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(256 * BasicBlock.expansion, num_classes)

        # 初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, out_channels, blocks, stride=1, norm_layer=None):
        norm_layer = norm_layer or nn.BatchNorm2d
        downsample = None
        if stride != 1 or self.inplanes != out_channels * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, out_channels * block.expansion, stride),
                norm_layer(out_channels * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, out_channels, stride, downsample, norm_layer=norm_layer))
        self.inplanes = out_channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, out_channels, norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x
    

class ResNet50Custom(nn.Module):
    def __init__(self, num_classes=idAmount):
        super(ResNet50Custom, self).__init__()
        self.inplanes = 32
        norm_layer = nn.BatchNorm2d

        # 为较小的灰度图设置合适的输入层
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = norm_layer(32)
        self.relu = nn.ReLU(inplace=True)

        # 设置模型各层
        self.layer1 = self._make_layer(Bottleneck, 32, 3, stride=1, norm_layer=norm_layer)
        self.layer2 = self._make_layer(Bottleneck, 64, 4, stride=2, norm_layer=norm_layer)
        self.layer3 = self._make_layer(Bottleneck, 128, 6, stride=2, norm_layer=norm_layer)
        self.layer4 = self._make_layer(Bottleneck, 256, 3, stride=2, norm_layer=norm_layer)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(256 * Bottleneck.expansion, num_classes)

        # 初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, out_channels, blocks, stride=1, norm_layer=None):
        norm_layer = norm_layer or nn.BatchNorm2d
        downsample = None
        if stride != 1 or self.inplanes != out_channels * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, out_channels * block.expansion, stride),
                norm_layer(out_channels * block.expansion)
            )

        layers = []
        layers.append(block(self.inplanes, out_channels, stride, downsample, norm_layer=norm_layer))
        self.inplanes = out_channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, out_channels, norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x
